import csv
import sqlite3
import pandas as pd
import utils.configFinRAG as configfinrag

def execute_sql_query(sql, cursor):
    """
    Executes an SQL query and returns the result and success flag.
    """
    try:
        cursor.execute(sql)
        result = cursor.fetchall()
        execution_result = str(result)
        success_flag = 1
    except Exception as e:
        execution_result = "error: " + str(e)
        success_flag = 0
        print(execution_result)
    return success_flag, execution_result

if __name__ == '__main__':
    # Load the SQL file containing questions and SQL statements
    question_sql_df = pd.read_csv(configfinrag.question_sql_path, delimiter=",", header=0)

    # Open the file to write SQL execution results
    with open(configfinrag.question_sql_check_path, 'w', newline='', encoding='utf-8-sig') as result_file:
        csv_writer = csv.writer(result_file)
        csv_writer.writerow(['问题id', '问题', 'SQL', '执行成功标志', '执行结果'])

        # Connect to the database
        db_connection = sqlite3.connect(r'C:\Users\李兴泽\Desktop\Fay-fay-agent-edition0830\Fay-fay-agent-edition0830\data\博金杯比赛数据.db')
        db_cursor = db_connection.cursor()

        # Execute SQL queries and write results to the file
        for index, row in question_sql_df.iterrows():
            if row['SQL'] != 'error':
                success_flag, execution_result = execute_sql_query(row['SQL'], db_cursor)
                csv_writer.writerow([row['问题id'], row['问题'], row['SQL'], success_flag, execution_result])
            else:
                print(f"Skipping SQL execution for question ID: {row['问题id']} due to previous error.")

        # Close the database connection
        db_connection.close()
